import torch
import torch.nn as nn
import numpy as np

# Focal Loss
# Multi-Label-CE


def build_weights(self, mode="default"):
    """
    class unbanlance weight
    """
    weights = np.ones(self.num_classes)
    for label, cnt in self.train_freq.items():
        weights[self.label_alphabet.get_index(label)] += cnt
    if mode == "reciprocal":
        self.weights = 1.0 / weights
    elif mode == "smooth":
        self.weights = 1.0 / (weights**0.06)
    else:
        self.weights = np.ones(self.num_classes)


def configure_loss(weights=[], type="bce"):

    weights = torch.tensor(weights).float().cuda()
    print(weights)

    if type == "bce":
        if len(weights):
            loss_fn = nn.BCEWithLogitsLoss(weight=weights)

        else:
            loss_fn = nn.BCEWithLogitsLoss()

    else:
        if len(weights):
            loss_fn = nn.CrossEntropyLoss(weight=weights)
        else:
            loss_fn = nn.CrossEntropyLoss()

    return loss_fn
